# -*- coding: utf-8 -*-

import os
import codecs
import pickle
import numpy as np
import pandas as pd
from keras.models import Sequential
from keras.models import Model
from keras.layers import Dense
from keras.layers.merge import concatenate
from keras.layers import LSTM
from keras.layers import Dropout
from keras.layers import CuDNNLSTM
from keras.layers import Input
from keras.layers import SpatialDropout1D
from keras.layers.convolutional import Conv1D
from keras.layers.convolutional import MaxPooling1D
from keras.layers.embeddings import Embedding
from keras.layers import Bidirectional
from keras.preprocessing.sequence import pad_sequences
from keras.preprocessing.text import Tokenizer
from keras.utils import np_utils
from keras.callbacks import EarlyStopping, ModelCheckpoint
from sklearn.model_selection import train_test_split

from util.attention import AttentionWithContext


def make_tokenizer(file_name, cache_file):
    if not os.path.exists(cache_file):
        tokenizer = Tokenizer(char_level=False)
        corpus = []
        with open(file_name) as reader:
            for line in reader:
                corpus.append(line.strip())
        tokenizer.fit_on_texts(corpus)
        with open(cache_file, "w") as writer:
            pickle.dump(tokenizer, writer, protocol=pickle.HIGHEST_PROTOCOL)
    else:
        tokenizer = pickle.load(open(cache_file))
    return tokenizer


def pad_sequence(text, tokenizer, sent_len):
    sequence = tokenizer.texts_to_sequences(text)
    return pad_sequences(sequence, maxlen=sent_len)


def make_embedding_matrix(tokenizer, vec_dim, embed_file, cache_file):
    if not os.path.exists(cache_file):
        embedding_index = {}
        with codecs.open(embed_file) as reader:
            next(reader)
            for line in reader:
                values = line.strip().split(" ")
                word = values[0]
                coefs = np.asarray(values[1:], dtype="float32")
                embedding_index[word] = coefs
        word_index = tokenizer.word_index
        embedding_matrix = np.zeros((len(word_index)+1, vec_dim))
        for word, index in word_index.items():
            vector = embedding_index.get(word)
            if vector is not None:
                embedding_matrix[index] = vector
        np.save(cache_file, embedding_matrix)
    else:
        embedding_matrix = np.load(cache_file)
    return embedding_matrix


def cnn_lstm_model(tokenizer, embedding_dim, embedding_matrix, max_length, num_class):
    # create the model
    model = Sequential()
    model.add(Embedding(len(tokenizer.word_index)+1, embedding_dim, weights=[embedding_matrix], input_length=max_length))
    model.add(SpatialDropout1D(0.3))
    # model.add(Conv1D(filters=16, kernel_size=3, padding='same', activation='relu'))
    # model.add(MaxPooling1D(pool_size=2))
    model.add(Bidirectional(CuDNNLSTM(50, return_sequences=True)))
    model.add(AttentionWithContext())
    model.add(Dropout(0.2))
    # model.add(Dense(32, kernel_initializer='glorot_uniform', activation='relu'))
    model.add(Dense(num_class, activation='softmax'))
    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    # model.summary()
    return model


def cnn_lstm_merge_model(app_tknz, app_embed_dim, app_embed_matrix, app_max_len,
                         brand_tknz, brand_embed_dim, brand_embed_matrix, brand_max_len,
                         num_class):
    app_input = Input(shape=(app_max_len,))
    app_part = Embedding(len(app_tknz.word_index)+1, app_embed_dim, weights=[app_embed_matrix])(app_input)
    app_part = SpatialDropout1D(0.3)(app_part)
    app_part = Conv1D(filters=16, kernel_size=3, padding='same', activation='relu')(app_part)
    app_part = MaxPooling1D(pool_size=2)(app_part)
    app_part = LSTM(50)(app_part)

    brand_input = Input(shape=(brand_max_len,))
    brand_part = Embedding(len(brand_tknz.word_index)+1, brand_embed_dim, weights=[brand_embed_matrix], trainable=False)(brand_input)
    # brand_part = SpatialDropout1D(0.2)(brand_part)
    # brand_part = Conv1D(filters=32, kernel_size=2, padding='same', activation='relu')(brand_part)
    # brand_part = MaxPooling1D(pool_size=2)(brand_part)
    brand_part = LSTM(5)(brand_part)

    merged = concatenate([app_part, brand_part])
    preds = Dense(num_class, activation='softmax')(merged)

    model = Model(inputs=[app_input, brand_input], outputs=preds)
    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    # model.summary()
    return model


if __name__ == '__main__':
    # app info
    appMaxLen = 174  # 174
    appEmbedDim = 50
    appTknz = make_tokenizer(os.path.join("data", "deviceid_packages_cooked.txt"), os.path.join("cache", "app.tokenizer"))
    appEmbeddings = make_embedding_matrix(appTknz, appEmbedDim, os.path.join("cache", "app_cbow_50d.vec"), os.path.join("cache", "app"))

    brandMaxLen = 10
    brandEmbedDim = 5
    brandTknz = make_tokenizer(os.path.join("data", "deviceid_brand_cooked.txt"), os.path.join("cache", "brand.tokenizer"))
    brandEmbeddings = make_embedding_matrix(brandTknz, brandEmbedDim, os.path.join("cache", "brand_cbow_5d.vec"), os.path.join("cache", "brand"))

    names = ["app", "brand", "label", "gender", "age"]
    # names = ["app", "label", "gender", "age"]
    df = pd.read_csv(os.path.join("data", "app_brand.train"), sep=",", names=names)
    print df.shape

    label = df.pop("label")
    y = np_utils.to_categorical(label)

    train_X, valid_X, train_y, valid_y = train_test_split(df, y, random_state=42, test_size=0.2, stratify=y)
    """
    train_input = [pad_sequence(train_X["app"], appTknz, appMaxLen), pad_sequence(train_X["brand"], brandTknz, brandMaxLen)]
    valid_input = [pad_sequence(valid_X["app"], appTknz, appMaxLen), pad_sequence(valid_X["brand"], brandTknz, brandMaxLen)]

    cnnLstm = cnn_lstm_merge_model(
        appTknz, appEmbedDim, appEmbeddings, appMaxLen,
        brandTknz, brandEmbedDim, brandEmbeddings, brandMaxLen,
        num_class=22
    )
    """
    train_input = [pad_sequence(train_X["app"], appTknz, appMaxLen)]
    valid_input = [pad_sequence(valid_X["app"], appTknz, appMaxLen)]

    cnnLstm = cnn_lstm_model(appTknz, appEmbedDim, appEmbeddings, appMaxLen, num_class=22)
    # """
    earlyStop = EarlyStopping(monitor="val_loss", patience=3, verbose=0)
    checkPoint = ModelCheckpoint("model/cnn_lstm_attention.h5", monitor="val_loss", save_best_only=True, verbose=0)

    cnnLstm.fit(train_input, train_y, validation_data=(valid_input, valid_y), epochs=15, batch_size=64, callbacks=[earlyStop, checkPoint])
# CUDA_VISIBLE_DEVICES="1,2,3" python lstm_model.py
